import json
from glob import glob
import random

LABLE_MAP = {
    "biaoqian_clear_correct": 0,
    "biaoqian_clear_wrong": 1,
    "biaoqian_blur": 2,
    "fangshuijiaobu": 3,
    "taoguan_zhiguan": 4,
    "taoguan_bowenguan": 5
}

CROSS_VALIDATION = {
    "train": 0.7,
    "val": 0.1,
    "test": 0.2
}

def collect_image(item, ind):
    image_id = {}
    image_id["coco_url"] = ""
    image_id["data_capture"] = ""
    image_id["flickr_url"] = ""
    image_id["file_name"] = item["imagePath"]
    image_id["height"] = item["imageHeight"]
    image_id["width"] = item["imageWidth"]
    image_id["id"] = ind

    return image_id

def collect_annotation(item, ind, anno_start_ind):
    item_anno = item["shapes"]
    anno_ids = []
    for i, bbox_item in enumerate(item_anno):
        anno_id = {}
        anno_id["id"] = anno_start_ind + i
        anno_id["image_id"] = ind
        anno_id["category_id"] = get_category_id(bbox_item["label"])
        anno_id["segmentation"], anno_id["bbox"] = points2bbox(bbox_item["points"])
        anno_id["iscrowd"] = 0
        anno_id["area"] = compute_area(bbox_item["points"])
        anno_ids.append(anno_id)
    return anno_ids

def collect_category():
    categories = []
    for key, value in LABLE_MAP.items():
        category = {}
        category["id"] = value
        category["name"] = key
        category["supercategory"] = "none"
        categories.append(category)
    return categories

def get_category_id(label):
    return LABLE_MAP[label]

def points2bbox(points):
    x1, y1 = points[0]
    x2, y2 = points[1]
    seg_anno = [x1, y1, x2, y1, x2, y2, x1, y2]
    seg_anno = [float(i) for i in seg_anno]
    bbox_anno = [x1, y1, (x2 - x1), (y2 - y1)]
    bbox_anno = [float(i) for i in bbox_anno]
    return [seg_anno], bbox_anno

def compute_area(points):
    x1, y1 = points[0]
    x2, y2 = points[1]
    return float((x2 - x1) * (y2 - y1))

if __name__ == '__main__':
    random.seed(42)
    
    data_dir = "/home/HDD/linzihao/datasets/liantong/CPRI/cpri/cpri/all_data"
    out_dir = "/home/HDD/linzihao/cpri_equipcheck/data_prepare/annotation_1"
    json_list = glob(data_dir + "/*.json")
    json_len = len(json_list)

    random.shuffle(json_list)
    modes = ['train', 'val', 'test']

    start_id = 0
    for mode in modes:
        anno_dict = {}
        anno_dict['images'] = []
        anno_dict['annotations'] = []
        anno_dict['categories'] = []

        anno_id = 0
        json_list_thismode = json_list[start_id : start_id + int(json_len * CROSS_VALIDATION[mode])]
        start_id += int(json_len * CROSS_VALIDATION[mode])
        for ind, json_file in enumerate(json_list_thismode):
            with open(json_file, 'r') as f:
                data = json.load(f)

            if "shapes" not in data.keys():
                continue

            anno_dict['images'].append(collect_image(data, ind))
            anno_dict['annotations'].extend(collect_annotation(data, ind, anno_id))
            anno_id += len(data["shapes"])

            anno_dict['categories'] = collect_category()

        with open(out_dir + f"/annotations_{mode}.json", 'w') as f:
            json.dump(anno_dict, f, indent=4)
        